from langchain_core.retrievers import BaseRetriever
from langchain_core.documents import Document
from langchain_core.callbacks.manager import CallbackManagerForRetrieverRun
from typing import List, Optional
from app.services.vector_store.milvus_service import MilvusService

class MilvusRetriever(BaseRetriever):
    """基于Milvus的检索器"""
    
    def __init__(self, milvus_service: MilvusService):
        """初始化检索器
        
        Args:
            milvus_service: Milvus服务实例
        """
        super().__init__()
        self.milvus_service = milvus_service
    
    def _get_relevant_documents(
        self, query: str, *, run_manager: Optional[CallbackManagerForRetrieverRun] = None
    ) -> List[Document]:
        """获取与查询相关的文档
        
        Args:
            query: 查询文本
            run_manager: 回调管理器
            
        Returns:
            相关文档列表
        """
        return self.milvus_service.similarity_search(query)